#coding:utf-8
#==============================================================================
#
"""
    Base class for datastore aware objects.
"""
#
#   Copyright (C) 2009 Google, Inc.
#
#   This file is part of ODB.
#
#   ODB is free software: you can redistribute it and/or modify it under the
#   terms of the GNU Lesser General Public License as published by the Free
#   Software Foundation, either version 3 of the License, or (at your option)
#   any later version.
#
#   ODB is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Lesser General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with ODB.  If not, see <http://www.gnu.org/licenses/>.
#
#==============================================================================

import re, urllib
import odb

_fixRx = re.compile(r'[\0|\x01]')
def _fixSegmentChar(match):
    return {'\0': '\x01\0', '\x01': '\x01\x01'}[match.group()]

def _encodeInteger(intVal):
    """
        Encode an integer (or long integer) value as a string that can be
        sorted lexically.
    """
    bytes = []
    neg = intVal < 0
    while intVal:
        bytes.append(chr(intVal % 254 + 2))
        intVal /= 254

        # gargghh, silly special case...
        if intVal == -1:
            break 

    # we're currently not going to deal with really big integers.  When we do,
    # special byte values of 255 and 2 will introduce a two byte length.  I
    # don't think we should ever support > 2 byte integer lengths.
    assert len(bytes) < 126

    if neg:
        size = chr(128 - len(bytes))
    else:
        size = chr(128 + len(bytes))
    bytes.append(size)
    bytes.reverse()
    return ''.join(bytes)

def _encodeSegment(segment):
    if segment is None:
        # encode None as an empty string.
        return ''
    elif isinstance(segment, (int, long)):
        return _encodeInteger(segment)
    elif isinstance(segment, unicode):
        return segment.encode('utf-8')
    else:
        return _fixRx.sub(_fixSegmentChar, segment)

def createKey(vals):
    """
        Create a key from a list of values.  Returns tuple<str, int> consisting of the
        initial key and the index of the beginning of the segment defined
        by the last range.

        parms:
            vals: [list<str or Matcher>] a list of values to convert to a key.
                Anything following the first Matcher object is ignored.
    """
    parts = []
    rangeSegmentIndex = -1
    for val in vals:
        if isinstance(val, Matcher):
            # join everything up to the range now so we can easily get the
            # length.
            if parts:
                parts = ['\0'.join(parts)]
                rangeSegmentIndex = len(parts[0]) + 1
            else:
                # this is the first segment, range segment starts at the
                # beginning.
                rangeSegmentIndex = 0
            parts.append(_encodeSegment(val.getStartSegment()))

            # quit now: everything following the first range is ignored.
            break
        else:
            parts.append(_encodeSegment(val))
    return '\0'.join(parts), rangeSegmentIndex

def _iterateOverKey(table, key, lastMatcher):
    # now we can finally do the query
    cursor = table.cursor()
    key, wildcardStart = createKey(key)
    keyBase = key[:wildcardStart]
    cursor.setToKey(key, exact = not lastMatcher)

    # do the search
    for k, val in cursor:

        # exact matches must exactly match
        if not lastMatcher:
            if k != key:
                break

        else:
            # ranges should always start with the common part of the key
            if not k.startswith(keyBase):
                break

            # last segment must fit the range.
            if lastMatcher.isOutside(k[wildcardStart:]):
                break

        yield k, val

class QueryParseError(Exception):
    """Raised when we fail to parse a query."""
    pass

_nameRx = re.compile(r'[a-z_]\w*', re.I)
_segmentRx = re.compile(r'[^/\[]*')
_filterRx = re.compile(r'[^\]]+')
_prefixRx = re.compile(r'([^:]+)\*$')
_intRx = re.compile(r'-?\d+$')

def _chomp(remaining, expr, error):
    m = expr.match(remaining)
    if not m:
        raise QueryParseError('%s: %s' % (error, repr(remaining)))
    return m.group(), remaining[m.end():]

def _unescape(text):
    """
        Unescapes the text, converting %XX sequences to binary values.
    """
    return urllib.unquote(text)

def _convertAtom(text):
    if _intRx.match(text):
        return long(text)
    else:
        return _unescape(text)

def _convertSegment(segment):

    # is it a complete wildcard?
    if segment == '*':
        return Wild()

    # is it a partial wildcard?
    m = _prefixRx.match(segment)
    if m:
        return Prefix(_unescape(m.group(1)))

    # embedded wildcards are verboten
    if '*' in segment:
        raise QueryParseError('Wildcard is only legal at the end of a '
                               'segment: %s' % repr(segment)
                              )

    # check for a range
    if ':' in segment:
        try:
            start, end = segment.split(':')
            start = _convertAtom(start) or None
            end = _convertAtom(end) or None
            return Range(start, end)
        except ValueError:
            raise QueryParseError('range must have exactly two parts: %s' %
                                   repr(segment)
                                  )


    return _convertAtom(segment)

def query(queryString, store = None):
    """
        Runs a string query against the datastore, iterates over the matching
        result set.  Raises QueryParseError if the string is malformed.

        A query string consists of a path consisting of
        -   a table name (and optional index)
        -   a slash and a slash separated list of key segments
        -   a python expression to be used as a filter on the matching results
            enclosed in square brackets

        This is intended for interactive use.  For queries within python
        programs, Model.select() is a much better tool.
    """
    r = queryString

    # get the table name
    table, r = _chomp(r, _nameRx, 'Table name not found')

    # check for an index
    if r.startswith('.'):
        r = r[1:]
        index, r = _chomp(r, _nameRx, 'Index name not found')
    else:
        index = None

    # check for a key path
    path = []
    gotWild = False
    lastMatcher = None
    while r and r.startswith('/'):
        r = r[1:]
        segment, r = _chomp(r, _segmentRx, 'Segment name expected')
        segment = _convertSegment(segment)

        if gotWild and (not isinstance(segment, Matcher) or
                        segment.isPartial()):
            raise QueryParseError('Non-wildcard or partial match follows '
                                   'wildcard in key path.')

        if isinstance(segment, Matcher):
            gotWild = True
            lastMatcher = segment
        path.append(segment)

    # if no path, make the entire path a wildcard
    if not path:
        path = [Wild()]
        lastMatcher = path[0]

    # check for a filter
    if r.startswith('['):
        r = r[1:]
        # XXX this needs to do python tokenization
        filter, r = _chomp(r, _filterRx, 'Malformed filter expression')
        filter = compile(filter, 'filter-string', 'eval')
    else:
        filter = None

    if not store:
        store = odb.getStore()

    table = store.getMap(table)
    if index:
        table = table.getMapIndex(index)

    if filter:
        for k, v in _iterateOverKey(table, path, lastMatcher):
            # XXX should use a lazy-wrapper for this
            globals = {}
            for attr in dir(v):
                globals[attr] = getattr(v, attr)
            if eval(filter, globals):
                yield k, v
    else:
        for k, v in _iterateOverKey(table, path, lastMatcher):
            yield k, v

class Matcher:
    """Base class of types that can match a range of keys."""

    def isPartial(self):
        """
            Returns true if the range represents a partial wildcard (i.e. it
            won't match anything).
        """
        raise NotImplementedError()

    def isOutside(self, segment):
        """
            Returns true if the key segment is outside of the range.

            parms:
                segment: [str] This is an _encoded_ segment - for example, if
                    the original segment was "test\0val", we'll get
                    "test\x01\x00val".
        """
        raise NotImplementedError()

    def getStartSegment(self):
        """
            Returns the first key segment that will be matched by the range (a
            string).
        """
        raise NotImplementedError()

class Wild(Matcher):
    """Matches the full remainder of any key."""

    def isPartial(eslf):
        return False

    def isOutside(self, segment):
        return False

    def getStartSegment(self):
        return ''

# since there is no additional state, this can be a constant.
WILD = Wild()

class Prefix(Matcher):
    """Matches any segment beginning with a specified prefix."""

    def __init__(self, prefix):
        self.__prefix = prefix
        self.__encoded = _encodeSegment(prefix)

    def isPartial(self):
        return False

    def isOutside(self, segment):
        return not segment.startswith(self.__encoded)

    def getStartSegment(self):
        return self.__prefix

class Range(Matcher):
    """Matches a range of values."""

    def __init__(self, start = None, end = None):
        self.__start = start
        if start is not None:
            self.__encodedStart = _encodeSegment(start)
        self.__end = end
        if end is not None:
            self.__encodedEnd = _encodeSegment(end)

    def isPartial(self):
        return True

    def isOutside(self, segment):
        return self.__start and segment < self.__encodedStart or \
               self.__end and segment > self.__encodedEnd

    def getStartSegment(self):
        return self.__start or ''

class Schema:
    """
        Stores the primary table and all indeces for a given model object.
    """

    def __init__(self, table, primKey, indeces = {}):
        self.table = table
        self.primKey = primKey
        self.indeces = indeces

    @staticmethod
    def createKey(obj, keySpec):
        """
            Creates a key for an object from the attributes that constitute
            its key specification.

            parms:
                obj: [Model] the object to create a key for.
                keySpec: [tuple<str>] a sequence of attributes to
                    join together into the key.
        """
        parts = []
        for attr in keySpec:
            val = getattr(obj, attr)
            try:
                parts.append(_encodeSegment(val))
            except TypeError, ex:
                raise TypeError('Encoding attribute %s: %s' % (attr, ex[0]))
        return '\0'.join(parts)

class KeyCollisionError(Exception):
    """
        Raised when we attempt to overwrite an existing key.
    """

    def __init__(self, tableName, key):
        """
            parms:
                tableName: [str] table or index name (for an index, this is
                    Table.Index) where the collision occurred.
                key: [str] the key that collided.
        """
        Exception.__init__(self,
                           'Key %s collided in table %s' %
                            (repr(key), tableName)
                           )
        self.tableName = tableName
        self.key = key

class Model:
    """
        Base class for objects that want to automatically store themselves in
        all of their indeces.
    """

    def __init__(self):
        # a dictionary mapping the index names to the full keys where the
        # object is stored for that index.  The primary key is stored under
        # the empty string:
        # { '': 'primary key value', 'Index1': 'Index1 key value' }
        # For unstored objects, this should be None.
        self.__keys = None

    def put(self, store=None):
        """
            Puts the object into its primary table and all of the indeces
            as defined in its schema.
        """
        if not store:
            store = odb.getStore()

        # start a transaction or sub-transaction
        txn = store.startTxn()
        newKeys = None
        try:

            # prepare to update the primary table
            table = store.getMap(self._schema.table)
            primKey = Schema.createKey(self, self._schema.primKey)
            orgKeys = self.__keys
            oldPrimKey = orgKeys and orgKeys['']

            # first create all of the new keys so we can store them with the
            # object.
            self.__keys = newKeys = {'': primKey}
            for indexName, indexKeySchema in self._schema.indeces.iteritems():
                self.__keys[indexName] = Schema.createKey(self, indexKeySchema)

            # if the primary key has changed, check for a collision on the new
            # key and delete the old key.
            primKeyChanged = oldPrimKey != primKey
            if primKeyChanged:
                if table.get(primKey):
                    raise KeyCollisionError(self._schema.table, primKey)
                if oldPrimKey is not None:
                    table.delete(oldPrimKey)

            table.put(primKey, self)

            # update all of the indeces
            for indexName in self._schema.indeces:
                index = table.getMapIndex(indexName)
                indexKey = newKeys[indexName]
                oldIndexKey = orgKeys and orgKeys[indexName]

                # if the index key changed, check for a collision with the new
                # one and delete the old one
                indexKeyChanged = oldIndexKey != indexKey
                if indexKeyChanged:
                    if index.get(indexKey):
                        raise KeyCollisionError('%s.%s' %
                                                (self._schema.table, indexName),
                                                indexKey
                                                )
                    if oldIndexKey is not None:
                        index.delete(oldIndexKey)

                # we only need to rewrite the index key if it changed or if
                # the primary key changed.
                if indexKeyChanged or primKeyChanged:
                    index.link(indexKey, primKey)

            txn.commit()
            txn = None
        finally:
            if txn:
                txn.abort()

                # restore the original keys
                if newKeys:
                    self.__keys = orgKeys

    def delete(self, store=None):
        """
            Remove the object from its primary table and all indeces.
        """
        if not store:
            store = odb.getStore()

        assert self.__keys, 'attempt to delete an unstored object'

        # start a transaction or sub-transaction
        txn = store.startTxn()
        try:
            # remove from the primary table
            table = store.getMap(self._schema.table)
            primKey = self.__keys['']
            table.delete(primKey)

            # delete from all of the indeces
            for indexName in self._schema.indeces:
                index = table.getMapIndex(indexName)
                try:
                    indexKey = self.__keys[indexName]
                except KeyError:
                    # __keys will not have an entry for the index if the
                    # index has been added since the object was stored
                    continue

                index.delete(indexKey)

            txn.commit()
            txn = None
            self.__keys = None
        finally:
            if txn:
                txn.abort()

    def __str__(self):
        attrs = self.__dict__.keys()
        attrs.sort()
        lines = ['%s.%s {' %
                  (self.__class__.__module__, self.__class__.__name__)
                 ]
        for attr in attrs:
            if not attr.startswith('_Model__'):
                lines.append('    %s: %s' % (attr, repr(getattr(self, attr))))
        lines.append('}')
        return '\n'.join(lines)

    @classmethod
    def select(cls, *parms, **kwargs):
        """
            Iterate over all records matching the specified criteria.

            The keyword and sequence arguments together define the contents of
            the keys to be matched: if there are sequence arguments, they are
            the values of key attributes of the table or index as defined
            in the schema.  If there are keyword arguments (except for the
            special "filter", "index" and "store" keyword arguments) they
            should correspond to attribute names defined in the key.

            Any undefined key attributes are wild-carded.  A portion of a key
            attribute may also be explicitly wild-carded by passing in a
            @Matcher instance.

            When specifying key attributes through keyword args, it is illegal
            to omit key attributes that occur prior to the last specified
            attribute in the key attribute order.  Likewise, it is illegal to
            specify a range object in any but the last specified attribute.
            For example, if our key is ('x', 'y'), both of the following
            are illegal:

                Obj.select(y = 'third')
                Obj.select(Prefix('foo'), 'second')

            parms:
                *parms: [any] Sequence args.  These are key parameter values
                    and may be of the type specified by the key.
                **kwargs: [any] Except for the subsequent reserved keyword
                    arguments, these may be used as an alternate means to
                    specify key attributes.
                index: [str] if specified, this is the name of the index to
                    use for the lookup.
                store: [odb.Store] the datastore object from which to load
                    tables.
                filter: [callable<Model>] a function to be called to filter
                    out a subset of the values.  Returns true if the object is
                    to be included in the result set, false if not.
        """
        store = None
        if 'store' in kwargs:
            store = kwargs['store']
            del kwargs['store']
        if 'index' in kwargs:
            index = kwargs['index']
            del kwargs['index']
        else:
            index = None

        if not store:
            store = odb.getStore()

        # get the table and the index, if specified.
        tableName = cls._schema.table
        table = store.getMap(tableName)
        keyAttrs = cls._schema.primKey
        if index:
            try:
                keyAttrs = cls._schema.indeces[index]
            except KeyError:
                raise ValueError('%s is not a valid index for %s objects' %
                                  (index, cls.__name__)
                                 )
            table = table.getMapIndex(index)
            tableName = '%s.%s' % (tableName, index)

        parmsLen = len(parms)
        if parmsLen > len(keyAttrs):
            raise ValueError('Too many arguments for key of %s, expected %s' %
                              (tableName, keyAttrs)
                             )

        # create a template key from the parameters - fill it with wildcards
        # so that we know if anything is missing.
        key = [Wild()] * len(keyAttrs)
        for i, val in enumerate(parms):
            key[i] = val

        # go through the remaining keyword arguments, verify that they are
        # part of the key and not already consumed by the sequence arguments
        for attr, val in kwargs.iteritems():
            try:
                attrIndex = list(keyAttrs).index(attr)
            except ValueError:
                raise ValueError('Attribute %s is not part of the key for %s' %
                                  (attr, tableName)
                                 )
            if attrIndex < parmsLen:
                raise ValueError('Attribute %s is already specified in '
                                  'sequence arguments.' % attr
                                 )

            # store the keyword parameter in the template key
            key[attrIndex] = val

        # verify that the key doesn't have any wildcards where it shouldn't
        # (after the last non-wildcard elem)
        gotWildcards = False
        keyIter = iter(key)
        lastMatcher = None
        for item in keyIter:
            itemIsRange = isinstance(item, Matcher)
            if gotWildcards and (not itemIsRange or item.isPartial()):
                raise ValueError('Non-wildcard or partial wildcard follows '
                                  'wildcard in key.'
                                 )
            elif not gotWildcards and itemIsRange:
                gotWildcards = True
                lastMatcher = item

        for item in _iterateOverKey(table, key, lastMatcher):
            yield item

    @classmethod
    def get(cls, *parms, **kwargs):
        """
            Returns a single object matching the key, returns None or raises
            a KeyError if there is no object matching the key (depending on
            whether the errorIfMissing flag is set.

            This is just like select() only designed for exact matches.  It
            does not preclude wildcard matches, however, if a wildcard is
            specified, it must match exactly one object, otherwise a
            ValueError will be raised.

            parms: see @select().parms
                errorIfMissing: [bool] if specified as a keyword argument,
                    this is true if the function should raise
        """

        # if we got the errorIfMissing arg, get its value and remove it so we
        # don't pass it on to select()
        if 'errorIfMissing' in kwargs:
            errorIfMissing = kwargs.get('errorIfMissing')
            del kwargs['errorIfMissing']
        else:
            errorIfMissing = False

        cursor = cls.select(*parms, **kwargs)
        try:
            key, result = cursor.next()
        except StopIteration:
            if errorIfMissing:
                raise KeyError(parms, kwargs)
            else:
                return None

        # make sure we get a StopIteration the next time
        try:
            cursor.next()
            raise ValueError('Multiple values matched')
        except StopIteration:
            return result
